#!/usr/bin/env python
# -*- conding: utf-8 -*-

"""
@Time     : 2024/10/29 6:43
@Author   : liujingmao
@File     : 1.条件边与循环构建工具调用Agent.py
"""

import dotenv
from langchain_community.tools import GoogleSerperRun
from langchain_community.tools.openai_dalle_image_generation import OpenAIDALLEImageGenerationTool
from langchain_community.utilities import GoogleSerperAPIWrapper
from langchain_community.utilities.dalle_image_generator import DallEAPIWrapper
from langchain_core.pydantic_v1 import BaseModel, Field
from langchain_openai import ChatOpenAI
from langgraph.checkpoint.memory import MemorySaver
from langgraph.prebuilt import create_react_agent

dotenv.load_dotenv()


class GoogleSerperArgsSchema(BaseModel):
    query: str = Field(description="执行谷歌搜索的查询语句")


class DallEArgsSchema(BaseModel):
    query: str = Field(description="输入应该是生成图像的文本提示(prompt)")


# 1.定义工具与工具列表
google_serper = GoogleSerperRun(
    name="google_serper",
    description=(
        "一个低成本的谷歌搜索API。"
        "当你需要回答有关时事的问题时，可以调用该工具。"
        "该工具的输入是搜索查询语句。"
    ),
    args_schema=GoogleSerperArgsSchema,
    api_wrapper=GoogleSerperAPIWrapper(),
)
dalle = OpenAIDALLEImageGenerationTool(
    name="openai_dalle",
    api_wrapper=DallEAPIWrapper(model="dall-e-3"),
    args_schema=DallEArgsSchema,
)
tools = [google_serper, dalle]

# 2.创建大语言模型
model = ChatOpenAI(model="gpt-4o-mini", temperature=0)

# 3.使用预构建的函数创建ReACT智能体
checkpointer = MemorySaver()
config = {"configurable": {"thread_id": 1}}
agent = create_react_agent(model=model, tools=tools, checkpointer=checkpointer)

# 4.调用智能体并输出内容
print(agent.invoke(
    {"messages": [("human", "你好，我叫茂茂，我喜欢游泳打篮球，你喜欢什么呢?")]},
    config=config,
))

# 5.二次调用检测图结构程序是否存在记忆
print(agent.invoke(
    {"messages": [("human", "你知道我叫什么吗?我喜欢什么")]},
))

"""
{'messages': [HumanMessage(content='你好，我叫慕小课，我喜欢游泳打球，你喜欢什么呢?', additional_kwargs={}, 
response_metadata={}, id='3685fb49-0497-4077-a4f4-0174b222d8a0'), 
AIMessage(content='你好，慕小课！我没有个人喜好，但我很喜欢帮助你解答问题和提供信息。
如果你有任何想聊的或者需要帮助的，随时告诉我！你最喜欢的运动是什么呢？', 
additional_kwargs={'refusal': None}, 
response_metadata={'token_usage': {'completion_tokens': 50, 'prompt_tokens': 171, 'total_tokens': 221, 
'completion_tokens_details': {'reasoning_tokens': 0}, 
'prompt_tokens_details': {'cached_tokens': 0}}, 
'model_name': 'gpt-4o-mini-2024-07-18', 
'system_fingerprint': 'fp_f59a81427f', 
'finish_reason': 'stop', 'logprobs': None}, 
id='run-866b088e-c455-41df-a691-adf4d8e79527-0', 
usage_metadata={'input_tokens': 171, 'output_tokens': 50, 'total_tokens': 221, 'input_token_details': {'cache_read': 0}, 'output_token_details': {'reasoning': 0}})]}
"""

"""
{'messages': [HumanMessage(content='你好，我叫茂茂，我喜欢游泳打篮球，你喜欢什么呢?', additional_kwargs={}, response_metadata={}, id='e4afccfd-c6b9-497b-9068-d386fcc8b100'), AIMessage(content='你好，茂茂！我没有个人喜好，但我很高兴听到你喜欢游泳和打篮球。这两项运动都很有趣且有益健康！你最喜欢的运动是什么时候进行的呢？', additional_kwargs={'refusal': None}, response_metadata={'token_usage': {'completion_tokens': 51, 'prompt_tokens': 171, 'total_tokens': 222, 'completion_tokens_details': {'reasoning_tokens': 0}, 'prompt_tokens_details': {'cached_tokens': 0}}, 'model_name': 'gpt-4o-mini-2024-07-18', 'system_fingerprint': 'fp_f59a81427f', 'finish_reason': 'stop', 'logprobs': None}, id='run-eed48a7a-1c32-4302-8138-c5c555eef7eb-0', usage_metadata={'input_tokens': 171, 'output_tokens': 51, 'total_tokens': 222, 'input_token_details': {'cache_read': 0}, 'output_token_details': {'reasoning': 0}})]}
"""
